﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using OpenCvSharp;
using OpenCvSharp.Dnn;

namespace YoloPredictor
{
    /// <summary>
    /// 适用于Darknet模型的Predictor
    /// <para>为项目FruitDetector编写</para>
    /// </summary>
    public class OpenCvDnn : IPredictor
    {
        private readonly string[] Labels;
        private readonly Net Network;
        public readonly PredictorConfig PerdiConfig;
        public double XRatio = 1, YRatio = 1;

        public OpenCvDnn(PredictorConfig conf)
        {
            Labels = File.ReadAllLines(conf.Names);
            Network = CvDnn.ReadNetFromDarknet(conf.Config, conf.Weights);
            Network.SetPreferableBackend(Backend.OPENCV);               //指定Darknet网络使用OpenCV
            Network.SetPreferableTarget(conf.UseCUDA ? Target.CUDA : Target.CPU);
            PerdiConfig = conf;
        }

        public List<Prediction> Predict(Mat picture, bool enableNMS = true)
        {
            XRatio = picture.Width / 416f;
            YRatio = picture.Height / 416f;

            var blob = CvDnn.BlobFromImage(picture, 1.0 / 255, new Size(416, 416), new Scalar(), true, false);
            Network.SetInput(blob);
            var outNames = Network.GetUnconnectedOutLayersNames();
            var outs = outNames.Select(_ => new Mat()).ToArray();       //为输出层准备Mat
            Network.Forward(outs, outNames);                            //正向运算，获得输出层结果
            return ReadResults(outs, picture, PerdiConfig.Threshold, PerdiConfig.NmsThreshold, enableNMS);
        }

        /// <summary>
        /// 获取检测结果
        /// </summary>
        /// <param name="output"></param>
        /// <param name="image"></param>
        /// <param name="threshold"></param>
        /// <param name="nmsThreshold"></param>
        /// <param name="nms"></param>
        /// <returns></returns>
        private List<Prediction> ReadResults(IEnumerable<Mat> output, Mat image, float threshold, float nmsThreshold, bool nms = true)
        {
            var result = new List<Prediction>();
            //for nms
            var classIds = new List<int>();
            var confidences = new List<float>();
            var probabilities = new List<float>();
            var boxes = new List<Rect2d>();

            var w = image.Width;
            var h = image.Height;
            /*
             YOLO3 标准输出结果
             0 1:中心点坐标                    2 3 : 宽、高
             4 : Confidence                  5 ~ 84 : Probability
            */
            const int skip = 5;   //跳过0~4

            foreach (var prob in output)
            {
                for (var i = 0; i < prob.Rows; i++)
                {
                    var confidence = prob.At<float>(i, 4);
                    if (confidence > threshold)
                    {
                        //计算命中可能性
                        Cv2.MinMaxLoc(prob.Row(i).ColRange(skip, prob.Cols), out _, out Point max);
                        var classId = max.X;
                        var probability = prob.At<float>(i, classId + skip);

                        if (probability > threshold)//非必要
                        {
                            var centerX = prob.At<float>(i, 0) * w;
                            var centerY = prob.At<float>(i, 1) * h;
                            var width = prob.At<float>(i, 2) * w;
                            var height = prob.At<float>(i, 3) * h;

                            if (!nms)
                            {
                                // result(无NMS)
                                result.Add(
                                    new Prediction()
                                    {
                                        ClassId = classId,
                                        ClassName = Labels[classId],
                                        Confidence = confidence,
                                        //Box = new Rect2d(centerX * XRatio, centerY * YRatio, width * XRatio, height * YRatio),
                                        Box = new Rect2d(centerX, centerY, width, height),
                                        Probability = probability
                                    }
                                );
                                continue;
                            }
                            classIds.Add(classId);
                            confidences.Add(confidence);
                            probabilities.Add(probability);
                            boxes.Add(new Rect2d(centerX, centerY, width, height));
                        }
                    }
                }
            }

            if (!nms) return result;

            //抑制重合检测框
            CvDnn.NMSBoxes(boxes, confidences, threshold, nmsThreshold, out int[] indices);
            foreach (var i in indices)
            {
                var box = boxes[i];
                result.Add(
                    new Prediction()
                    {
                        ClassId = classIds[i],
                        ClassName = Labels[classIds[i]],
                        Confidence = confidences[i],
                        //Box = new Rect2d(box.X * XRatio, box.Y * YRatio, box.Width * XRatio, box.Height * YRatio),
                        Box = box,
                        Probability = probabilities[i]
                    }
                );
            }
            return result;
        }
    }
}
